1213G - Path Queries - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
 
 
#include <ext/pb_ds/assoc_container.hpp> // Common file
#include <ext/pb_ds/tree_policy.hpp>
#include <functional> // for less
using namespace __gnu_pbds;
 

typedef tree<int, null_type,less_equal<int>, rb_tree_tag,
             tree_order_statistics_node_update>
    ordered_set;
    
    
typedef long double LD;
typedef long long ll;
#define pb push_back
#define pob pop_back
#define mp make_pair
#define rep(n) for (ll i = 0; i < n; i++)
#define FOR(i,a,b) for (i = a; i < b; i++)
#define repd(n) for (ll i = n-1; i >= 0; i--)
#define FORD(i,a,b) for (i = a; i >= b; i--)
#define remax(a,b) a = max(a,b)
#define remin(a,b) a = min(a,b)
#define all(v) v.begin(),v.end()
#define back(v) v.rbegin(),v.rend()
#define wl(t) while(t --)
#define inrange(i,a,b) ((i>=min(a,b)) && (i<=max(a,b)))
#define isodd %2==1
#define iseven %2==0
typedef map<ll,ll> mii;
typedef vector<ll> vi;
typedef vector<vi> vvi;
typedef pair<ll,ll> pii;
typedef vector<pii> vpii;
#define take(a,n) rep(n)cin>>a[i];
#define give(a,n) rep(n)cout<<a[i]<<' ';
 
#define FLSH fflush(stdout)
#define fileIO(name) \
    freopen(name".in", "r", stdin); \
    freopen(name".out", "w", stdout);
#define PRECISION(x) cout << fixed << setprecision(x); 
#define FAST_IO ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
 
const ll MOD = 1000000007;
const ll FMOD = 998244353;
const double PI=3.1415926;
 
template<typename T> T gcd(T a, T b){return(b?__gcd(a,b):a);}
template<typename T> T lcm(T a, T b){return(a*(b/gcd(a,b)));}
int add(int a, int b, int c = MOD){int res=a+b;return(res>=c?res-c:res);}
int mod_neg(int a, int b, int c = MOD){int res;if(abs(a-b)<c)res=a-b;else res=(a-b)%c;return(res<0?res+c:res);}
int mul(int a, int b, int c = MOD){ll res=(ll)a*b;return(res>=c?res%c:res);}
int muln(int a, int b, int c = MOD){ll res=(ll)a*b;return ((res%c)+c)%c;}
ll mulmod(ll a,ll b, ll m = MOD){ll q = (ll)(((LD)a*(LD)b)/(LD)m);ll r=a*b-q*m;if(r>m)r%=m;if(r<0)r+=m;return r;}
template<typename T>T expo(T e, T n){T x=1,p=e;while(n){if(n&1)x=x*p;p=p*p;n>>=1;}return x;}
template<typename T>T power(T e, T n, T m = MOD){T x=1,p=e;while(n){if(n&1)x=mul(x,p,m);p=mul(p,p,m);n>>=1;}return x;}
#define SORT(v) sort(all(v))
#define REVERSE(v) reverse(all(v))
#define sz(a) ll((a).size())
#define tr(container, it) for (decltype (container.begin()) it = container.begin(); it != container.end(); it++)
#define cpresent(container, element)(find(all(container), element) != container.end())
#define present(container, element)(container.find(element) != container.end())

ll fac[1000000];


ll power(ll x, ll y, ll p)
{
 
    // Initialize answer
    ll res = 1;
    x=x%p; 
    // Check till the number becomes zero
    while (y > 0) {
 
        // If y is odd, multiply x with result
        if (y % 2 == 1)
            res = (res * x)%p;
 
        // y = y/2
        y = y >> 1;
 
        // Change x to x^2
        x = (x * x)%p;
    }
    return res % p;
}


ll modInverse(ll n,ll p)
{
    return power(n, p - 2, p);
}

ll nCr(ll n, ll r, ll p)
{
    // If n<r, then nCr should return 0
    if(r<0)return 0;
    if (n < r)
        return 0;
    // Base case
    if (r == 0)
        return 1;
    return (((fac[n] * modInverse(fac[r], p)) % p)
            * modInverse(fac[n - r], p)) % p;
}



// ll count(vector<ll> a, ll m, ll n)
// {
//     ll odd = 0, even = 0;
//     ll counter, i, j, p = 1;
//     ll pow_set_size = (1 << n);
// //  Given N prime numbers and a number M, find out how many numbers 
// //  from [1 to M] are divisible by any of the N given prime numbers. 
//     for (counter = 1; counter < pow_set_size; counter++) {
//         p = 1;
//         for (j = 0; j < n; j++) {
//             if (counter & (1 << j)) {
//                 p *= a[j];
//             }
//         }
//         if (__builtin_popcount(counter) & 1)odd += (m / p);
//         else even += (m / p);
//     }
 
//     return odd - even;
// }

// vi a,t1,t2;

// void build1(int v, int tl, int tr) {
//     if (tl == tr) {
//         t1[v] = a[tl];
//     } else {
//         int tm = (tl + tr) / 2;
//         build1(v*2, tl, tm);
//         build1( v*2+1, tm+1, tr);
//         t1[v] = gcd(t1[v*2] ,t1[v*2+1]);
//     }
// }
// ll q1(int v, int tl, int tr, int l, int r) {
//     if (l > r) 
//         return 0;
//     if (l == tl && r == tr) {
//         return t1[v];
//     }
//     int tm = (tl + tr) / 2;
//     return gcd(q1(v*2, tl, tm, l, min(r, tm)),q1(v*2+1, tm+1, tr, max(l, tm+1), r));
// }

// ll N=3167;
// ll N = 31627;
// ll N=3000000;
// vector<bool> isprime(N+5, true);
// vi pr;
 
// void seive(){
//     isprime[0] = isprime[1] = false;
//     for (ll i = 2; i * i <= N; i++) {
//         if (isprime[i]) {
//             for (ll j = i * i; j <= N; j += i)
//                 isprime[j] = false;
//         }
//     }
//     for (ll i = 2; i <= N; ++i)
//         if(isprime[i]) pr.pb(i);
// }

// vector<vi> adj;
// vi v;
// void dfs(ll i,ll p){
//     v.pb(i);
//     for(auto it:adj[i]){
//         if(it==p)continue;
//         dfs(it,i);
        
//     }
// }
ll N=300000;
vi parent(N),siz(N);

void make_set(int v) {
    parent[v] = v;
    siz[v] = 1;
}

int find_set(int v) {
    if (v == parent[v])
        return v;
    return parent[v] = find_set(parent[v]);
}

void union_sets(int a, int b) {
    a = find_set(a);
    b = find_set(b);
    if (a != b) {
        if (siz[a] < siz[b])
            swap(a, b);
        parent[b] = a;
        siz[a] += siz[b];
    }
}


int main(void)
{
	FAST_IO;
	ll m,k,q,e,f,o,x,t,g,j,h,y,z,p,n,mn,mx,u,r,l;
	t=1;
	
	fac[0] = 1;
    for (ll i = 1; i < 10000; i++)fac[i] = (fac[i - 1] * i) % FMOD;
    
    
    
    // seive();
    // cin>>t;
    while(t--){
        
        // a.clear();
        // a.resize(n);
        // t1.clear();
        // t1.resize(n*4);
        // vi a(n,0);
        // take(a,n);
        cin>>n>>m;
        // adj.resize(n);
        // rep(n)adj[i].clear();
        // vi a(n),b(n),c(n);
        // v.clear();
        // v.resize(n,0);
        ll a[n-1],b[n-1];
        pii v[n-1];
        rep(n-1){
            cin>>a[i]>>b[i]>>r;;
            a[i]--;
            b[i]--;
            v[i]={r,i};
        }
        sort(v,v+n-1);
        j=0;
        rep(n)make_set(i);
        vi ans;
        rep(200000+1){
            if(i==0){
                ans.pb(i);
                continue;
            }
            h=ans[i-1];
            while(v[j].first<=i && j<n-1){
                k=v[j].second;
                j++;
                h+=siz[find_set(a[k])]*siz[find_set(b[k])];
                union_sets(a[k],b[k]);
                
            }
            ans.pb(h);
        }
        rep(m){
            cin>>r;
            cout<<ans[r]<<" ";
        }
        
    }
	   

	
    return 0;
}